import torch
import torch.nn as nn
import torch.nn.functional as F
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.InstanceNorm3d(out_channels),
            nn.LeakyReLU(inplace=True),
            # nn.ReLU(inplace=True),
            nn.Dropout3d(p=0.1),
            nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.InstanceNorm3d(out_channels),
            nn.LeakyReLU(inplace=True),
            nn.Dropout3d(p=0.2) # 修复版为p=0.1
        )

    def forward(self, x):
        return self.double_conv(x)

class Down(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool3d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)

class Up(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.up = nn.ConvTranspose3d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        
        diffZ = x2.size()[2] - x1.size()[2]
        diffY = x2.size()[3] - x1.size()[3]
        diffX = x2.size()[4] - x1.size()[4]
        
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2,
                        diffZ // 2, diffZ - diffZ // 2])
        
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

class UNet3D(nn.Module):
    def __init__(self, n_channels=3, n_classes=1):
        super(UNet3D, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        
        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        
        self.down4 = DoubleConv(512, 1024)
        
        self.up1 = Up(1024, 512)
        self.up2 = Up(512, 256)
        self.up3 = Up(256, 128)
        self.up4 = Up(128, 64)

        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)        # -> [B, 64, 64, 64, 64]
        x2 = self.down1(x1)     # -> [B, 128, 32, 32, 32]
        x3 = self.down2(x2)     # -> [B, 256, 16, 16, 16]
        x4 = self.down3(x3)     # -> [B, 512, 8, 8, 8]
        
        x5 = self.down4(x4)     # -> [B, 1024, 8, 8, 8]
        
        x = self.up1(x5, x4)    # -> [B, 512, 16, 16, 16]
        x = self.up2(x, x3)     # -> [B, 256, 32, 32, 32]
        x = self.up3(x, x2)     # -> [B, 128, 64, 64, 64]
        x = self.up4(x, x1)     # -> [B, 64, 64, 64, 64]

        logits = self.outc(x)   # -> [B, 1, 64, 64, 64]
        return logits

if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device} - {torch.cuda.get_device_name(0)}")
    input_tensor = torch.randn(2, 1, 64, 64, 64).to(device)
    
    model = UNet3D(n_channels=1, n_classes=1).to(device)
    output = model(input_tensor)
    
    print("输入维度:", input_tensor.shape)
    print("输出维度:", output.shape)  # 应为 [2, 1, 64, 64, 64]

